// Copyright 2022 Chen Jun
#include "armor_detector/pnp_solver.hpp"

#include <opencv2/calib3d.hpp>
#include <vector>

#include "rclcpp/rclcpp.hpp"

namespace rm_auto_aim
{
PnPSolver::PnPSolver(
  const std::array<double, 9> & camera_matrix, const std::vector<double> & dist_coeffs,
  const rclcpp::NodeOptions & options)
: Node("pnp_solver", options),
  camera_matrix_(cv::Mat(3, 3, CV_64F, const_cast<double *>(camera_matrix.data())).clone()),
  dist_coeffs_(cv::Mat(1, 5, CV_64F, const_cast<double *>(dist_coeffs.data())).clone())
{
  // Unit: m
  //计算大装甲板和小装甲板的半宽和半高；/ 1000.0是为了将单位从毫米转换为米
  constexpr double small_half_y = SMALL_ARMOR_WIDTH / 2.0 / 1000.0;
  constexpr double small_half_z = SMALL_ARMOR_HEIGHT / 2.0 / 1000.0;
  constexpr double large_half_y = LARGE_ARMOR_WIDTH / 2.0 / 1000.0;
  constexpr double large_half_z = LARGE_ARMOR_HEIGHT / 2.0 / 1000.0;

  // Start from bottom left in clockwise order
  // Model coordinate: x forward, y left, z up
  //构建物体坐标系下的大装甲板和小装甲板的四个角点；物体坐标系原点设置在装甲板的中心（三维物体点集）
  small_armor_points_.emplace_back(cv::Point3f(0, small_half_y, -small_half_z));
  small_armor_points_.emplace_back(cv::Point3f(0, small_half_y, small_half_z));
  small_armor_points_.emplace_back(cv::Point3f(0, -small_half_y, small_half_z));
  small_armor_points_.emplace_back(cv::Point3f(0, -small_half_y, -small_half_z));

  large_armor_points_.emplace_back(cv::Point3f(0, large_half_y, -large_half_z));
  large_armor_points_.emplace_back(cv::Point3f(0, large_half_y, large_half_z));
  large_armor_points_.emplace_back(cv::Point3f(0, -large_half_y, large_half_z));
  large_armor_points_.emplace_back(cv::Point3f(0, -large_half_y, -large_half_z));
}
/*const Armor & armor：表示装甲板的对象，包含了装甲板的相关信息。
cv::Mat & rvec：输出参数，用于存储计算得到的旋转向量。
cv::Mat & tvec：输出参数，用于存储计算得到的平移向量。*/
bool PnPSolver::solvePnP(const Armor & armor, cv::Mat & rvec, cv::Mat & tvec)
{
  //创建用于存储装甲板在图像中的四个角点的变量image_armor_points（二维图像点集）
  std::vector<cv::Point2f> image_armor_points;

  // Fill in image points
  image_armor_points.emplace_back(armor.left_light.bottom);
  image_armor_points.emplace_back(armor.left_light.top);
  image_armor_points.emplace_back(armor.right_light.top);
  image_armor_points.emplace_back(armor.right_light.bottom);

  // Solve pnp
  auto object_points = armor.type == ArmorType::SMALL ? small_armor_points_ : large_armor_points_;
  //调用cv::solvePnP函数计算装甲板的旋转向量和平移向量
  /*
  代码调用OpenCV的cv::solvePnP函数来解决姿态估计问题。cv::solvePnP函数接受以下参数：
  object_points：三维物体点集。
  image_armor_points：二维图像点集。
  camera_matrix_：相机内参矩阵。
  dist_coeffs_：相机畸变系数。
  rvec：输出的旋转向量。（输出值）
  tvec：输出的平移向量。（输出值）——平移向量的z分量 就是 物体距离相机的深度信息
  false：表示不使用初始估计。
  cv::SOLVEPNP_IPPE：使用IPPE算法进行求解。
  */
  // 调用 cv::solvePnP 函数计算装甲板的旋转向量和平移向量
  bool success = cv::solvePnP(
    object_points, image_armor_points, camera_matrix_, dist_coeffs_, rvec, tvec, false,
    cv::SOLVEPNP_IPPE);

  // 如果成功，输出目标距离
  if (success) {
    RCLCPP_INFO(this->get_logger(), "PnP result--Object distance: %.2f meters", tvec.at<double>(2));
  } else {
    RCLCPP_ERROR(this->get_logger(), "PnP solve failed!");
  }

  return success;
}
/*
calculateDistanceToCenter函数于计算图像中的某个点到相机中心的距离。
输入参数：一个cv::Point2f类型的参数image_point，表示图像中的某个点。

*/
float PnPSolver::calculateDistanceToCenter(const cv::Point2f & image_point)
{
  /*提取出相机的光心坐标。*/
  //获取相机内参矩阵中第0行第2列的值，即相机光心在x轴上的坐标cx
  float cx = camera_matrix_.at<double>(0, 2);
  //获取相机内参矩阵中第1行第2列的值，即相机光心在y轴上的坐标cy
  float cy = camera_matrix_.at<double>(1, 2);
  //计算图像中的某个点到相机中心的欧几里德距离
  return cv::norm(image_point - cv::Point2f(cx, cy));
}

}  // namespace rm_auto_aim
